import numpy as np
from cvfo_mask import CvFo
from tqdm import tqdm
import paddle
import pandas as pd

block_size=8
voc=pd.read_pickle("voc.pandas_pickle")
data=pd.read_pickle("dataset.pandas_pickle")
model=CvFo(len(voc),512,6,block_size)
# model.load_dict(paddle.load("cvfo_model.pdparams"))
# model.load_dict(paddle.load("cvfo_model.pdparams"))
opt=paddle.optimizer.Adam(learning_rate=0.00005, parameters=model.parameters())
loss_func= paddle.nn.CrossEntropyLoss()
epochs=800
batch_size=500
epoch=list(range(0,len(data),batch_size))
bar=tqdm(epoch*epochs)
for i in bar:
    if i==0:
        np.random.shuffle(data)
    batch_data=paddle.to_tensor(data[i:i+batch_size]).astype('int64')
    labels=batch_data[:,-1:,1:]
    out=model(batch_data[:,:-1,:-1])


    loss=loss_func(out,labels.squeeze(1))
    bar.set_description("loss:%.4f"%loss.item())
    opt.clear_grad()
    loss.backward()
    opt.step()


    if i==epoch[-1]:
        paddle.save(model.state_dict(),"cvfo_model_rank.pdparams")
paddle.save(model.state_dict(),"cvfo_model_rank.pdparams")